Now Loading ...
-
FNet: Mixing Tokens with Fourier Transforms
논문 정보
제목: FNet: Mixing Tokens with Fourier Transforms
arXiv: 2105.03824
Introduction
트랜스포머 아키텍처의 핵심은 attention 메커니즘이다. attention은 하나의 토큰을 표현할 때 입력 시퀀스의 다른 모든 토큰과의 연관성을 계산하고, 이를 가중치로 정보를 종합한다. 이는 “모든 토큰이 서로를 참조할 수 있어야 한다”는 귀납적 편향(inductive bias)이 포함된 것이다. 따라서 레이어를 거칠수록 입력 토큰들의 조합으로부터 점점 더 상위 수준의 표현이 형성되며, 이러한 구조 덕분에 트랜스포머가 다양한 문법적(syntactic), 의미적(semantic) 관계를 포착할 수 있었다.
이 논문에서는 트랜스포머 인코더의 비교적 복잡한 self-attention 레이어를, 토큰을 섞는(token mixing) 더 단순한 메커니즘으로 완전히 대체할 수 있는지를 검증한다.
먼저 attention 서브레이어를 학습 가능한 두 개의 행렬곱으로 교체하였다. 하나는 시퀀스 차원을, 다른 하나는 은닉(hidden) 차원을 섞는다. 입력 $X \in \mathbb{R}^{N \times d}$에 대해 이는 다음과 같이 표현된다.
\[Y = W_{\text{seq}} \, X \, W_{\text{h}}, \qquad W_{\text{seq}} \in \mathbb{R}^{N \times N}, \; W_{\text{h}} \in \mathbb{R}^{d \times d}\]
이러한 단순한 선형 혼합만으로도 유의미한 성능이 나오자, 저자들은 더 빠르게 계산할 수 있는 구조화된(structured) 선형 변환으로 범위를 넓혀 실험하였다.
그 결과, 학습 파라미터가 전혀 없는 푸리에 변환(Fourier Transform)이 학습 가능한 선형 혼합과 거의 같은 성능을 보이는 것을 확인했으며, $O(N \log N)$의 FFT알고리즘 덕분에 긴 입력에 대해서도 매우 효율적으로 동작하였다. 이 모델을 FNet이라 명명하여 소개한다.
이 연구는 attention을 효율적으로 개선하려는 “efficient Transformer” 연구들과도 맞닿아 있다. attention 행렬을 희소화하거나 선형화하여 몇몇 연구들은 $O(N)$ 수준의 복잡도를 달성한다. 그러나 이들 기법은 근사 과정에서 발생하는 큰 상수 비용을 숨기고 있어, 실제로는 FNet보다 확장성이 떨어지는 경우가 많다.
이 논문에서는 크게 다음 4가지 결과를 보여준다:
단순한 선형 변환만으로도 텍스트의 다양한 관계를 모델링할 수 있다. 파라미터가 없는 푸리에 변환에 feed-forward 레이어의 MLP를 결합하는 것만으로도 충분한 성능이 나온다. 이는 적어도 일부 NLP 과제에서 attention이 트랜스포머 성능의 핵심 요소가 아닐 수도 있음을 시사한다.
FNet은 속도, 메모리, 정확도 사이에서 좋은 절충점을 제공한다. GLUE 벤치마크에서 BERT-Base 정확도의 92%, BERT-Large 정확도의 97%를 달성하면서, 학습 속도는 GPU에서 80%, TPU에서 70% 빠르다.
소수의 attention 레이어만 추가해도 정확도를 대부분 회복할 수 있다. 마지막 두 레이어만 self-attention으로 되돌린 FNet-Hybrid는 BERT 정확도의 97~99%를 달성하면서도 40~70% 빠르다. 이는 attention이 정확도 향상에 도움은 되지만, 모든 레이어에 필요하지 않을 수 있음을 보여준다.
FNet은 긴 입력에서 efficient Transformer보다 더 나은 속도–정확도 향상을 보인다. Long Range Arena(LRA) 벤치마크에서 FNet은 가장 정확한 efficient Transformer들과 비슷한 정확도를 보이면서, GPU를 이용한 모든 시퀀스 길이에서의 학습, 추론이 더 빨랐다.
Model
Discrete Fourier Transform
푸리에 변환은 함수를 그것을 구성하는 주파수 성분들로 분해하여 표현한다. 길이 $N$인 시퀀스 ${x_n}$ ($n = 0, \dots, N-1$)이 주어졌을 때, 이산 푸리에 변환(DFT)은 다음과 같이 정의된다.
\[X_k = \sum_{n=0}^{N-1} x_n \, e^{-\frac{2\pi i}{N} nk}, \qquad 0 \le k \le N-1\]
각 출력 $X_k$는 모든 입력 $x_n$에 회전 인자(twiddle factor) $e^{-\frac{2\pi i}{N}nk}$를 곱해 더한 값이다. 즉 출력의 각 성분이 시퀀스 전체의 정보를 담게 되며, FNet은 바로 이 성질을 토큰 혼합(token mixing)에 활용한다. 또한 회전 인자가 위치 인덱스 $n$에 의존하므로, 변환 결과에는 토큰의 위치 정보도 함께 반영된다.
즉 Fourier 서브레이어는 $[N, D]$ 입력을 고정된 무늬들로 측정하여, 출력 행렬의 각 원소가 모든 토큰과 모든 특징의 정보를 담도록 재표현하는 역할까지만 하는 것이다. 섞인 표현에서 의미 있는 패턴을 추출하는 일은 뒤따르는 feed-forward 서브레이어가 학습을 통해 수행한다.
예시를 통해 원리를 살펴보면 다음과 같다.
입력과 DFT 행렬
토큰 4개, 임베딩 차원 4인 입력 $X \in \mathbb{R}^{4 \times 4}$를 생각해 보자. 행은 토큰, 열은 임베딩 차원(특징)이다.
\[X = \begin{array}{c|cccc}
& h_0 & h_1 & h_2 & h_3 \\ \hline
\text{나는} & 1 & 0 & 2 & 1 \\
\text{밥을} & 0 & 1 & 1 & 0 \\
\text{정말} & 2 & 1 & 0 & 1 \\
\text{먹었다} & 1 & 0 & 1 & 2
\end{array}\]
$N = 4$일 때 회전 인자는 $e^{-\frac{2\pi i}{4}nk} = (-i)^{nk}$이므로, DFT 행렬은 다음과 같다. (정규화 계수는 생략)
\[W = \begin{bmatrix}
1 & 1 & 1 & 1 \\
1 & -i & -1 & i \\
1 & -1 & 1 & -1 \\
1 & i & -1 & -i
\end{bmatrix}\]
$W$의 각 행은 서로 다른 빠르기로 진동하는 기본 패턴이다.
$k$
패턴 (위치 0~3)
의미
0
$[1,\ 1,\ 1,\ 1]$
변화 없음 (전체 합)
1
$[1,\ -i,\ -1,\ i]$
길이 4 동안 한 번 진동
2
$[1,\ -1,\ 1,\ -1]$
한 칸마다 부호가 바뀜 (가장 빠른 진동)
3
$[1,\ i,\ -1,\ -i]$
$k=1$의 켤레 (반대 방향 회전)
$W$는 대칭 행렬이므로, 두 1D DFT는 행렬곱으로 다음과 같이 표현된다.
$\mathcal{F}_{\text{h}}(X) = XW$: 오른쪽에서 곱하며, 각 행(토큰) 안에서 특징들을 섞는다.
$\mathcal{F}_{\text{seq}}(X) = WX$: 왼쪽에서 곱하며, 각 열(특징) 안에서 토큰들을 섞는다.
따라서 Fourier 서브레이어는 다음과 같이 쓸 수 있다.
\[Y = \Re\big(\mathcal{F}_{\text{seq}}(\mathcal{F}_{\text{h}}(X))\big) = \Re(WXW)\]
이는 학습 가능한 선형 혼합 \(W_{\text{seq}} X W_{\text{h}}\)에서 두 학습되던 행렬을 고정된 DFT 행렬로 바꾼 형태이다.
1단계: 은닉 차원 DFT (각 토큰 안에서 섞기)
각 행에 DFT를 적용한다. 예를 들어 “나는” 행 $[1, 0, 2, 1]$은 다음과 같이 계산된다.
\[\begin{aligned}
k=0&: \; 1 + 0 + 2 + 1 = 4 \\
k=1&: \; 1 + 0\cdot(-i) + 2\cdot(-1) + 1\cdot(i) = -1 + i \\
k=2&: \; 1 - 0 + 2 - 1 = 2 \\
k=3&: \; 1 + 0\cdot(i) + 2\cdot(-1) + 1\cdot(-i) = -1 - i
\end{aligned}\]
네 행에 모두 적용하면 다음과 같다.
\[XW = \begin{array}{c|cccc}
& k_h{=}0 & k_h{=}1 & k_h{=}2 & k_h{=}3 \\ \hline
\text{나는} & 4 & -1+i & 2 & -1-i \\
\text{밥을} & 2 & -1-i & 0 & -1+i \\
\text{정말} & 4 & 2 & 0 & 2 \\
\text{먹었다} & 4 & 2i & 0 & -2i
\end{array}\]
이 단계에서 행은 여전히 토큰이다. 각 토큰은 자기 자신의 특징들만 섞었을 뿐, 다른 토큰의 정보는 아직 들어오지 않았다. 반면 열은 더 이상 원래의 특징 $h_m$이 아니라 특징 축의 주파수 성분 $k_h$가 된다.
2단계: 시퀀스 DFT (토큰끼리 섞기)
이번에는 각 열에 DFT를 적용한다. 첫 번째 열 $[4, 2, 4, 4]^\top$은 다음과 같이 계산된다.
\[\begin{aligned}
k=0&: \; 4 + 2 + 4 + 4 = 14 \\
k=1&: \; 4 + 2\cdot(-i) + 4\cdot(-1) + 4\cdot(i) = 2i \\
k=2&: \; 4 - 2 + 4 - 4 = 2 \\
k=3&: \; 4 + 2\cdot(i) + 4\cdot(-1) + 4\cdot(-i) = -2i
\end{aligned}\]
네 열에 모두 적용하면 다음과 같다.
\[WXW = \begin{array}{c|cccc}
& k_h{=}0 & k_h{=}1 & k_h{=}2 & k_h{=}3 \\ \hline
k_{\text{seq}}{=}0 & 14 & 2i & 2 & -2i \\
k_{\text{seq}}{=}1 & 2i & -6+2i & 2 & 0 \\
k_{\text{seq}}{=}2 & 2 & 2 & 2 & 2 \\
k_{\text{seq}}{=}3 & -2i & 0 & 2 & -6-2i
\end{array}\]
3단계: 실수부 추출
마지막으로 실수부만 취하면 Fourier 서브레이어의 출력이 된다.
\[Y = \Re(WXW) = \begin{bmatrix}
14 & 0 & 2 & 0 \\
0 & -6 & 2 & 0 \\
2 & 2 & 2 & 2 \\
0 & 0 & 2 & -6
\end{bmatrix}\]
두 DFT의 적용 순서는 결과에 영향을 주지 않는다. 행렬곱의 결합법칙에 의해 $W(XW) = (WX)W$이기 때문이다.
출력의 의미
행렬곱 $WXW$의 $(a, b)$ 원소를 풀어 쓰면 다음과 같다.
\[Y_{ab} = \Re\Big(\sum_{n,m} W_{an}\, X_{nm}\, W_{mb}\Big) = \sum_{n,m} X_{nm}\, P^{(a,b)}_{nm}, \qquad P^{(a,b)}_{nm} = \cos\!\left(\tfrac{\pi}{2}(an + bm)\right)\]
즉 $Y_{ab}$는 입력 $X$와 고정된 2차원 물결 무늬 $P^{(a,b)}$를 원소별로 곱해 모두 더한 값이다. 다시 말해 $X$가 그 무늬와 얼마나 닮았는지를 측정한 값이다. 예를 들어 $P^{(0,0)}$은 모든 원소가 1인 무늬이므로, $Y_{00} = 14$는 모든 토큰과 모든 특징을 더한 문장 전체의 총량이 된다.
이 관점에서 $Y$의 행과 열은 다음과 같이 해석할 수 있다.
행은 특정 토큰이 아니라, 문장 전체를 바라보는 하나의 관점이다. 행 $a$는 네 토큰을 패턴 $a$로 묶은 결과이다. 예를 들어 행 0은 네 토큰을 그대로 더한 문장 요약이고, 행 2는 $[+, -, +, -]$ 패턴으로 인접 토큰 간의 차이를 묶은 것이다.
열은 원래 특징이 아니라, 특징들을 조합해 만든 파생 특징이다. 열 $b$는 각 토큰의 특징을 패턴 $b$로 요약한 값이 토큰을 따라 어떻게 변하는지를 나타낸다.
결과적으로 출력의 모든 원소는 입력의 16개 값 전부로부터 계산되며, 이것이 Fourier 서브레이어가 토큰을 섞는 방식이다.
해석 시 유의할 점
1. 의미적 해석이 아니다.
$Y$의 각 원소에 문법적 관계와 같은 언어적 의미가 있는 것은 아니다. 무늬 $P^{(a,b)}$는 입력과 무관하게 고정되어 있기 때문이다. Attention은 무엇을 얼마나 섞을지를 입력이 결정하는 반면, FNet은 섞는 방식이 고정되어 있고 섞인 결과를 어떻게 해석할지만 이후 레이어에서 학습한다.
2. 실수부만 취하면 정보가 일부 중복된다.
$k = 1$과 $k = 3$은 서로 켤레 관계이므로 $Y_{ab} = Y_{(4-a)\bmod 4,\,(4-b)\bmod 4}$가 성립한다. 실제로 $Y_{11} = Y_{33} = -6$이다. 따라서 16개의 값이 모두 독립적이지는 않으며, 이 때문에 실수부만 취한 변환은 역변환이 불가능하다.
3. residual 연결이 토큰 정체성을 되살린다.
Fourier 서브레이어 뒤에는 $\text{LayerNorm}(X + Y)$가 이어진다. 따라서 $n$번째 행은 다음 두 정보를 함께 가진다.
\[\underbrace{X_{n,:}}_{\text{토큰 } n \text{ 자신의 특징}} + \underbrace{Y_{n,:}}_{\text{문장 전체를 } n\text{번 패턴으로 본 요약}}\]
$Y$ 단독으로는 행이 토큰에 대응하지 않지만, residual 덕분에 각 행은 “자기 토큰의 정보 + 전역 정보”를 갖게 된다. 이후 feed-forward 서브레이어는 각 행을 독립적으로 처리하는 연산임에도, 이미 모든 토큰의 정보가 섞인 입력을 받으므로 사실상 문장 전체를 보게 된다.
4. 비선형성은 feed-forward 레이어에만 있다.
Fourier 서브레이어 또한 선형 변환이다.
토큰 혼합이 없다면
토큰 혼합의 중요성은 Fourier 서브레이어를 완전히 제거한 FF-only 모델과 비교하면 분명해진다. 논문에서 FF-only 모델은 이진 분류 과제에서 약 50%, 즉 무작위 추측 수준의 정확도에 그쳤다. 또한 시퀀스 방향 DFT만 남기고 은닉 방향 DFT를 제거한 모델은 정확도가 다소 떨어졌지만 FF-only보다는 훨씬 좋은 성능을 보였다. 은닉 방향 혼합은 feed-forward 레이어도 수행하는 일이므로, Fourier 서브레이어에서 핵심적인 역할은 토큰 방향의 혼합임을 알 수 있다.
FNet architecture
그림과 같이 FNet은 attention이 없는(attention-free) 트랜스포머 인코더 아키텍처이다. 각 인코더 블록은 Fourier 서브레이어 → Add & Normalize → Feed Forward → Add & Normalize 순으로 구성되며, 이 블록을 $N$개 쌓은 뒤 dense 레이어로 최종 출력을 만든다. 즉 기존 트랜스포머 인코더 블록에서 self-attention 서브레이어만 Fourier 서브레이어로 교체한 구조이다.
Fourier 서브레이어는 (시퀀스 길이, 은닉 차원) 크기의 임베딩 입력에 2D DFT를 적용한다. 2D DFT는 시퀀스 차원을 따라 수행하는 1D DFT $\mathcal{F}{\text{seq}}$와 은닉 차원을 따라 수행하는 1D DFT $\mathcal{F}{\text{h}}$를 차례로 적용한 것과 같으며, 출력에서는 실수부만 취한다.
\[y = \Re\big(\mathcal{F}_{\text{seq}}(\mathcal{F}_{\text{h}}(x))\big)\]
두 1D DFT는 교환 가능하므로 적용 순서는 결과에 영향을 주지 않는다. 또한 실수부만 남기기 때문에, 뒤따르는 feed-forward 서브레이어나 출력 레이어가 복소수를 다룰 수 있도록 수정할 필요가 없다. Fourier 서브레이어 자체에는 학습 파라미터가 없다.
이 변환의 효과는 앞의 예시에서 보았듯이, 출력의 각 원소가 모든 토큰의 정보를 담게 된다는 것이다. 따라서 각 위치를 독립적으로 처리하는 feed-forward 서브레이어가 시퀀스 전체의 정보를 충분히 얻을 수 있다.
시간 영역과 주파수 영역의 교대
푸리에 변환에는 쌍대성(duality)이 있어, DFT를 두 번 적용하면 원래 신호가 (순서만 뒤집힌 채) 상수배되어 돌아온다.
\[\mathcal{F}(\mathcal{F}(x))_n = N \cdot x_{(-n) \bmod N}\]
즉 정방향 DFT를 한 번 더 적용하는 것은 역변환과 사실상 같은 역할을 한다. 따라서 FNet의 블록을 차례로 통과하는 과정은, 입력을 “시간” 영역(토큰 위치)과 주파수 영역 사이에서 번갈아 오가게 하는 것으로 볼 수 있다.
이때 합성곱 정리(convolution theorem)에 의해, 주파수 영역에서 계수를 곱하는 것은 시간 영역에서 합성곱을 하는 것과 같다.
\[\mathcal{F}^{-1}\big(\mathcal{F}(x) \odot \mathcal{F}(g)\big) = x \circledast g\]
따라서 주파수 영역에 있는 표현에 feed-forward 레이어의 가중치를 곱하는 것은, 원래 영역에서 보면 그와 관련된 커널로 합성곱을 하는 것에 해당한다. 이런 관점에서 저자들은 FNet을 곱셈과 합성곱을 번갈아 수행하는 모델로 해석할 수 있다고 설명한다. 다만 이는 직관적인 해석이며, 실제로는 실수부 추출, residual 연결, 비선형 활성화 함수가 끼어 있으므로 엄밀한 등가 관계는 아니다.
임베딩
DFT의 회전 인자 $e^{-\frac{2\pi i}{N}nk}$가 위치 인덱스 $n$에 의존하므로, Fourier 서브레이어는 그 자체로 위치 정보를 반영한다. 따라서 FNet은 원리적으로 positional embedding이 필요 없다. 그럼에도 논문에서는 BERT와 깔끔한 비교를 위해, BERT와 동일하게 단어 임베딩, 절대 위치 임베딩, 문장 구분(type) 임베딩을 합한 입력 임베딩을 사용하였다.
Implementation
실험 결과, GPU에서는 실험한 모든 시퀀스 길이(512–8192 토큰)에서 FFT가 행렬곱보다 빨랐다. 반면 TPU에서는 비교적 짧은 시퀀스(4096 토큰 이하)의 경우, DFT 행렬을 미리 계산해 캐싱해 두고 행렬곱으로 DFT를 계산하는 편이 FFT보다 빨랐다. 그보다 긴 시퀀스에서는 TPU에서도 FFT가 더 빨랐다.
이에 따라 GPU 구현에서는 항상 FFT를 사용하고, TPU 구현에서는 길이 4096까지는 행렬곱으로, 그보다 긴 시퀀스에서는 FFT로 2D DFT를 계산한다.
저자들은 이러한 차이의 원인으로 두 가지를 든다.
TPU는 GPU보다도 행렬곱에 더 고도로 최적화된 하드웨어이다.
GPU의 FFT 구현이 TPU의 FFT 구현보다 효율적이다.
따라서 TPU의 FFT 구현이 개선된다면 TPU에서의 FNet 성능도 더 향상될 것으로 예상된다. 모델은 JAX와 Flax 프레임워크로 구현되었다.
Results
1. Transfer learning (GLUE)
BERT와 동일한 Base/Large 설정으로 C4 데이터셋에서 사전학습한 뒤, GLUE에서 fine-tuning하여 비교하였다.
FNet은 BERT-Base 정확도의 92%, BERT-Large 정확도의 97%를 달성하면서 GPU에서 약 80%, TPU에서 약 70% 빠르게 학습된다. Large에서 격차가 줄어든 것은 FNet-Large가 BERT-Large보다 학습이 안정적이었기 때문으로 보인다.
FF-only는 이진 분류 과제에서 50% 수준에 그쳐 학습에 실패했다. 즉 토큰 혼합은 필수적이다.
Random의 낮은 성능은 아무렇게나 섞는 것이 아니라 구조화된 혼합이 필요함을 보여준다.
FNet-Hybrid는 마지막 두 Fourier 서브레이어만 self-attention으로 교체한 모델로, 적은 속도 손실로 BERT 정확도의 97~99%를 회복한다.
Linear는 Base에서 FNet보다 약간 정확하지만, 메모리를 더 많이 쓰고 긴 시퀀스로의 확장성이 떨어지며, Large에서는 학습이 불안정했다.
2. 속도–정확도 비교
모델 크기를 다양하게 바꿔 가며 비교하면, 큰 모델에서는 BERT와 FNet-Hybrid가, 작은 모델에서는 FNet과 Linear가 가장 좋은 속도–정확도 절충(Pareto frontier)을 보인다. 즉 FNet은 자원이 제한된 환경의 경량 모델에서 특히 유리하다.
3. Long Range Arena (LRA)
정확도: FNet의 평균 정확도(55.30)는 가장 정확한 모델인 Transformer(55.83)와 비슷하며, 대부분의 efficient Transformer보다 높다.
속도: GPU에서는 모든 시퀀스 길이(512–8192)에서 비교 모델 중 가장 빠르다. 길이 2048에서 Transformer 대비 학습은 3.2배, 추론은 3.4배 빠르다.
메모리: 모든 길이에서 메모리 사용량이 가장 적으며, Transformer가 메모리 부족(OOM)으로 학습하지 못하는 길이 8192에서도 7.4GB만 사용한다.
Deep Learning
· 2026-10-27
-
Andrej Karpathy - Makemore Part 4: Becoming a Backprop Ninja
원본 영상 및 자료
Youtube
Github
해당 영상에서는 역전파를 단순히 loss.backward()로 호출하는 것이 아닌, 직접 gradient를 하나하나 정의해보며 backprop의 동작을 구현한다.
Starter Code
이전과 동일한 MLP코드 위에서 진행.
block_size = 3
n_embd = 10 # the dimensionality of the character embedding vectors
n_hidden = 64 # the number of neurons in the hidden layer of the MLP
g = torch.Generator().manual_seed(2147483647) # for reproducibility
C = torch.randn((vocab_size, n_embd), generator=g)
# Layer 1
W1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * (5/3)/((n_embd * block_size)**0.5)
b1 = torch.randn(n_hidden, generator=g) * 0.1 # using b1 just for fun
# Layer 2
W2 = torch.randn((n_hidden, vocab_size), generator=g) * 0.1
b2 = torch.randn(vocab_size, generator=g) * 0.1
# BatchNorm parameters
bngain = torch.randn((1, n_hidden))*0.1 + 1.0
bnbias = torch.randn((1, n_hidden))*0.1
parameters = [C, W1, b1, W2, b2, bngain, bnbias]
직접 gradient를 계산하고 PyTorch gradient와 동일한지 확인하기 위한 utility function이다.
def cmp(s, dt, t):
ex = torch.all(dt == t.grad).item()
app = torch.allclose(dt, t.grad)
maxdiff = (dt - t.grad).abs().max().item()
print(f'{s:15s} | exact: {str(ex):5s} | approximate: {str(app):5s} | maxdiff: {maxdiff}')
loss 는 cross entropy인 것은 동일한데, 메서드 하나로 선언한 것이 아닌, 각각의 연산들을 분리하여 변수로 선언하였다. 아래의 변수들의 관계를 따라가며 각각의 단계에서의 gradient를 직접 구해볼 것이다.
emb = C[Xb] # embed the characters into vectors
embcat = emb.view(emb.shape[0], -1) # concatenate the vectors
# Linear layer 1
hprebn = embcat @ W1 + b1 # hidden layer pre-activation
# BatchNorm layer
bnmeani = 1/n*hprebn.sum(0, keepdim=True)
bndiff = hprebn - bnmeani
bndiff2 = bndiff**2
bnvar = 1/(n-1)*(bndiff2).sum(0, keepdim=True)
bnvar_inv = (bnvar + 1e-5)**-0.5
bnraw = bndiff * bnvar_inv
hpreact = bngain * bnraw + bnbias
# Non-linearity
h = torch.tanh(hpreact) # hidden layer
# Linear layer 2
logits = h @ W2 + b2 # output layer
# cross entropy loss (same as F.cross_entropy(logits, Yb))
logit_maxes = logits.max(1, keepdim=True).values
norm_logits = logits - logit_maxes # subtract max for stability
counts = norm_logits.exp()
counts_sum = counts.sum(1, keepdims=True)
counts_sum_inv = counts_sum**-1
probs = counts * counts_sum_inv
logprobs = probs.log()
loss = -logprobs[range(n), Yb].mean()
1. Backproping the atomic compute graph
아래의 코드는 각각의 gradient를 전부 구한 코드이다.
dlogprobs = torch.zeros_like(logprobs)
dlogprobs[range(n), Yb] = -1.0/n
dprobs = (1.0 / probs) * dlogprobs
dcounts_sum_inv = (counts * dprobs).sum(1, keepdim=True)
dcounts = counts_sum_inv * dprobs
dcounts_sum = (-counts_sum**-2) * dcounts_sum_inv
dcounts += torch.ones_like(counts) * dcounts_sum
dnorm_logits = counts * dcounts
dlogits = dnorm_logits.clone()
dlogit_maxes = (-dnorm_logits).sum(1, keepdim=True)
dlogits += F.one_hot(logits.max(1).indices, num_classes=logits.shape[1]) * dlogit_maxes
dh = dlogits @ W2.T
dW2 = h.T @ dlogits
db2 = dlogits.sum(0)
dhpreact = (1.0 - h**2) * dh
dbngain = (bnraw * dhpreact).sum(0, keepdim=True)
dbnraw = bngain * dhpreact
dbnbias = dhpreact.sum(0, keepdim=True)
dbndiff = bnvar_inv * dbnraw
dbnvar_inv = (bndiff * dbnraw).sum(0, keepdim=True)
dbnvar = (-0.5*(bnvar + 1e-5)**-1.5) * dbnvar_inv
dbndiff2 = (1.0/(n-1))*torch.ones_like(bndiff2) * dbnvar
dbndiff += (2*bndiff) * dbndiff2
dhprebn = dbndiff.clone()
dbnmeani = (-dbndiff).sum(0)
dhprebn += 1.0/n * (torch.ones_like(hprebn) * dbnmeani)
dembcat = dhprebn @ W1.T
dW1 = embcat.T @ dhprebn
db1 = dhprebn.sum(0)
demb = dembcat.view(emb.shape)
dC = torch.zeros_like(C)
for k in range(Xb.shape[0]):
for j in range(Xb.shape[1]):
ix = Xb[k,j]
dC[ix] += demb[k,j]
기본 개념
코드에서 d로 시작하는 모든 변수는 loss를 그 변수로 편미분한 값이다.
\[\texttt{dprobs} = \frac{\partial \text{loss}}{\partial \texttt{probs}}, \qquad \texttt{dW2} = \frac{\partial \text{loss}}{\partial \texttt{W2}}\]
기준은 항상 최종 스칼라인 loss로 고정되어 있다. 중간 변수끼리의 미분(예: probs를 counts_sum_inv로 미분한 것)은 국소 미분(local derivative) 이며, 별도의 이름을 붙이지 않고 그 자리에서 곧바로 곱해 사용한다.
이 방식의 중요한 성질은, loss가 스칼라이므로, 스칼라를 텐서로 미분한 결과는 원래 텐서와 shape이 정확히 같다는 것이다.
dprobs.shape == probs.shape # 항상 참
dW2.shape == W2.shape # 항상 참
수식을 잘못 유도했는지 확인하는 1차 검증 도구 shape이 안 맞으면 무조건 틀린 것이고, 반대로 shape이 맞는 조합이 하나뿐인 경우에는 그것이 곧 정답인 경우가 많다.
연쇄법칙: 국소 미분 × upstream gradient
역전파의 모든 줄은 예외 없이 다음 형태다.
\[\underbrace{\frac{\partial \text{loss}}{\partial x}}_{\texttt{dx (구하려는 것)}} = \underbrace{\frac{\partial \text{loss}}{\partial y}}_{\texttt{dy (이미 알고 있음)}} \times \underbrace{\frac{\partial y}{\partial x}}_{\text{국소 미분}}\]
여기서 dy를 upstream gradient라고 부른다. 이미 뒤쪽에서 계산되어 흘러들어온 값이다. 우리가 각 단계에서 실제로 해야 할 일은 국소 미분을 구해서 곱하는 것뿐이고, 그래서 계산 그래프를 출력에서 입력 방향(역순)으로 올라가야 한다.
예를 들어 forward가
probs = counts * counts_sum_inv
일 때, counts_sum_inv에 대한 국소 미분은 $\partial\,\texttt{probs} / \partial\,\texttt{counts_sum_inv} = \texttt{counts}$ 이므로
dcounts_sum_inv = counts * dprobs # 국소 미분 × upstream
이 된다. (실제 코드에는 .sum(1, keepdim=True)가 더 붙으며, 자세한 내용은 이후 내용에서 다룸)
규칙 A - 여러 번 쓰인 변수는 gradient를 더한다
forward에서 한 변수가 여러 곳에 쓰였다면, loss로 가는 경로가 여러 개라는 뜻이다. 다변수 연쇄법칙에 따라 각 경로에서 흘러온 gradient를 모두 합산해야 한다.
counts가 대표적인 예다.
counts_sum = counts.sum(1, keepdims=True) # 경로 1
counts_sum_inv = counts_sum**-1
probs = counts * counts_sum_inv # 경로 2
그래서 dcounts는 두 번에 걸쳐 채워진다.
dcounts = counts_sum_inv * dprobs # 경로 2에서의 gradient
...
dcounts += torch.ones_like(counts) * dcounts_sum # 경로 1에서의 gradient
이 때문에 역전파 순서가 중요해진다. 계산 그래프의 위상 정렬(topological order)의 역순을 지켜야 한다.
규칙 B - forward의 sum과 backward의 broadcast
가장 자주 쓰이고, 알아두면 절반 이상이 자동으로 풀리는 규칙이다.
forward에 sum이 있으면 → backward에는 broadcast
forward에 broadcast(복제)가 있으면 → backward에는 sum
forward broadcast –> backward sum
probs = counts * counts_sum_inv
# (n, V) (n, 1) ← broadcast
counts_sum_inv의 값들이 열 방향(axis=1)으로 브로드캐스팅 되어 V개의 열로 복사되어 원소별로 곱해졌다. 즉 forward에서 같은 값이 V번 복제되어 쓰인 것이고, 이건 “여러 번 쓰인 변수”의 경우와 동일하게 V개 경로에서 온 gradient를 전부 더해야 한다.
dcounts_sum_inv = (counts * dprobs).sum(1, keepdim=True)
shape으로도 확인된다. counts * dprobs는 (n, V)인데 dcounts_sum_inv는 (n, 1)이어야 하므로, axis=1을 없애는 sum이 필요하다.
같은 논리가 그대로 적용되는 곳들:
db2 = dlogits.sum(0) # b2가 n개 행에 broadcast됨
dbngain = (bnraw * dhpreact).sum(0, keepdim=True)
dbnbias = dhpreact.sum(0, keepdim=True)
dbnvar_inv = (bndiff * dbnraw).sum(0, keepdim=True)
dlogit_maxes = (-dnorm_logits).sum(1, keepdim=True) # 국소미분이 -1
dbnmeani = (-dbndiff).sum(0)
keepdim은 shape을 유지해야 한다면 사용한다.
forward sum –> backward broadcast
반대의 상황이다.
counts_sum = counts.sum(1, keepdims=True)
# (n, V) → (n, 1)
counts_sum의 한 원소는 그 행의 모든 counts 원소에 대해 국소 미분이 1이다. 따라서 dcounts_sum (n,1)을 (n,V)로 그대로 펼쳐서(복제해서) 전달하면 된다.
dcounts += torch.ones_like(counts) * dcounts_sum
sum에 상수 배가 붙어 있으면 그 상수도 함께 전달한다.
# bnvar = 1/(n-1) * bndiff2.sum(0, keepdim=True)
dbndiff2 = (1.0/(n-1)) * torch.ones_like(bndiff2) * dbnvar
# bnmeani = 1/n * hprebn.sum(0, keepdim=True) 평균은 sum의 상수배로 볼 수 있음
dhprebn += 1.0/n * (torch.ones_like(hprebn) * dbnmeani)
sum과 broadcast는 서로의 전치(transpose) 관계인 선형 변환이기 때문에 이러한 규칙이 생긴 것이다.
broadcast는 $\mathbb{R}^1 \to \mathbb{R}^V$ 로 가는 선형사상이고, 그 행렬은 1로 채워진 열벡터다.
sum은 $\mathbb{R}^V \to \mathbb{R}^1$ 로 가는 선형사상이고, 그 행렬은 1로 채워진 행벡터다.
역전파는 본질적으로 야코비안의 행렬의 전치를 곱하는 연산(vector-Jacobian product) 이므로, forward에서 한쪽을 썼다면 backward에서는 그 전치인 다른 쪽이 나온 것이다.
연산 유형별 정리
(1) 인덱싱으로 값 선택
loss = -logprobs[range(n), Yb].mean()
logprobs는 (n, V)지만, loss 계산에 실제로 쓰인 건 각 행에서 정답에 해당하는 위치 하나씩이다. 쓰이지 않은 값은 loss에 아무 영향을 주지 않으므로 gradient가 0이다.
\[\text{loss} = -\frac{1}{n}\sum_{i} \texttt{logprobs}[i, Yb_i] \;\Longrightarrow\; \frac{\partial \text{loss}}{\partial \texttt{logprobs}[i, Yb_i]} = -\frac{1}{n}\]
dlogprobs = torch.zeros_like(logprobs) # 기본값은 0
dlogprobs[range(n), Yb] = -1.0/n # 뽑힌 자리에만 값을 넣음
뽑힌 위치로만 gradient를 되돌려 보낸다.
(2) 원소별 연산
log, exp, tanh처럼 원소별로 동작하는 함수는 국소 미분도 원소별이므로, 그냥 원소별 곱만 하면 된다.
dprobs = (1.0 / probs) * dlogprobs # log'(x) = 1/x
dnorm_logits = counts * dcounts # exp'(x) = exp(x) = counts
dhpreact = (1.0 - h**2) * dh # tanh'(x) = 1 - tanh(x)^2
dcounts_sum = (-counts_sum**-2) * dcounts_sum_inv # (x^-1)' = -x^-2
dbnvar = (-0.5*(bnvar + 1e-5)**-1.5) * dbnvar_inv # (x^-0.5)' = -0.5x^-1.5
(3) max
logit_maxes = logits.max(1, keepdim=True).values
norm_logits = logits - logit_maxes
max는 여러 값 중 하나만 골라서 통과시킨다. 따라서 gradient도 최댓값이었던 그 자리로만 전달되고 나머지는 0이다.
dlogit_maxes = (-dnorm_logits).sum(1, keepdim=True)
dlogits += F.one_hot(logits.max(1).indices, num_classes=logits.shape[1]) * dlogit_maxes
여기서 dlogit_maxes를 출력해보면 값이 사실상 0이다. softmax는 입력 전체를 같은 값만큼 평행이동해도 결과가 변하지 않으므로, 어떤 값을 빼든 loss가 달라지지 않기 때문이다. 수치 안정성을 위한 이 중간 과정이 loss에 영향을 주지 않는다는 걸 gradient가 직접 보여주는 case이다.
(4) 행렬곱
logits = h @ W2 + b2
# (n,V) = (n,H) @ (H,V) + (V,)
dh = dlogits @ W2.T # (n,V) @ (V,H) = (n,H)
dW2 = h.T @ dlogits # (H,n) @ (n,V) = (H,V)
db2 = dlogits.sum(0) # b2는 n개 행에 broadcast --> sum (규칙 B)
하나하나 인덱스를 풀어서 유도할 수도 있지만, shape 비교로 빠르게 답을 얻을 수 있다. dh는 h와 같은 (n,H)여야 하고, 있는 것은 dlogits (n,V)와 W2 (H,V)뿐이다. 이 둘로 (n,H)를 만드는 방법은 dlogits @ W2.T 하나뿐이다. dW2도 마찬가지로 (H,V)를 만드는 조합이 h.T @ dlogits 하나뿐이다.
(5) 임베딩 조회
emb = C[Xb]
(1)의 인덱싱으로 값 선택을 행 단위로 한 것이다. C의 한 행(하나의 문자 임베딩)은 배치 안에서 여러 번 참조될 수 있으므로, 규칙 A에 따라 그 모든 사용처의 gradient를 더해야 한다.
dC = torch.zeros_like(C)
for k in range(Xb.shape[0]):
for j in range(Xb.shape[1]):
ix = Xb[k,j]
dC[ix] += demb[k,j] # 인덱스에 해당하는 행에 gradient전파
이중 루프는 느리므로 실제로는 dC.index_add_(0, Xb.view(-1), demb.view(-1, C.shape[1])) 같은 형태로 쓴다.
요약
forward 연산
backward 연산
sum (평균 포함)
broadcast (상수배 유지)
broadcast / 복제
해당 축으로 sum
변수를 여러 곳에 사용
각 경로의 gradient를 +=로 누적
원소별 함수 $f$
$f’$ 를 원소별 곱
max
argmax 위치로 전달
인덱싱 (gather)
해당하는 원소에만 gradient전파
행렬곱 A @ B
dA = dOut @ B.T, dB = A.T @ dOut
지금 구하려는 d변수의 shape은 원래 변수와 같아야 한다.
국소 미분을 구해 upstream gradient와 곱한다.
forward에서 값이 복제되었거나 여러 번 쓰였다면, backward에서는 그만큼 합산한다.
2. cross entropy를 한 번에 미분하기
Exercise 1에서 계산 그래프를 노드 단위로 따라가며 gradient를 계산했다면, 이번에는 forward를 F.cross_entropy 한번으로 묶어 중간 노드를 거치지 않고 gradient를 loss --> logits로 바로 전달하는 것이다.
# before: 8개의 원자 연산
logit_maxes = logits.max(1, keepdim=True).values
norm_logits = logits - logit_maxes
counts = norm_logits.exp()
counts_sum = counts.sum(1, keepdims=True)
counts_sum_inv = counts_sum**-1
probs = counts * counts_sum_inv
logprobs = probs.log()
loss = -logprobs[range(n), Yb].mean()
# now: 한 줄
loss_fast = F.cross_entropy(logits, Yb)
\[\frac{\partial L}{\partial \ell_{ij}} = \frac{1}{n}\left(p_{ij} - \delta_{j, y_i}\right) \qquad\Longrightarrow\qquad \texttt{dlogits} = \frac{1}{n}\big(\texttt{probs} - \texttt{onehot}(Y)\big)\]
핵심은 $-\log(e^{\ell_y} / \sum_k e^{\ell_k}) = -\ell_y + \log\sum_k e^{\ell_k}$ 로 $\log$와 $\exp$가 상쇄된다는 점이다. 그 결과 Exercise 1에서 거쳐야 했던 max 라우팅, 역수의 미분, 브로드캐스팅 합산이 하나도 남지 않는다.
dlogits = F.softmax(logits, 1) # probs
dlogits[range(n), Yb] -= 1 # - onehot(Y)
dlogits /= n # / n
결과
$n$을 곱해 평균을 되돌려보면 softmax 확률과 정확히 같고, 정답 위치에서만 1이 빠져 있다.
F.softmax(logits, 1)[0] # [0.03, 0.11, 0.07, ...]
dlogits[0] * n # [0.03, 0.11, 0.07, ...] 인데 정답 자리만 -0.93
위치
gradient
경사하강 시 logit 변화
정답 $y$
$p_y - 1 < 0$
증가
나머지 $j$
$p_j > 0$
감소
크기에도 의미가 있다. 오답인데 확률을 많이 준 쪽일수록 더 세게 눌리고, 정답에 이미 확신이 있으면 $p_y \approx 1$ 이라 거의 건드리지 않는다. 틀린 만큼만 고친다는 것이 식에 그대로 들어 있다.
그리고 각 행의 합은 0이다. 행합이 0이라는 건 logits 한 행 전체를 같은 값만큼 밀어도 loss가 변하지 않는다는 뜻이다.
dlogits[0].sum() # 1e-9
Exercise 1에서 dlogit_maxes가 0이었던 것 - softmax는 입력 전체를 평행이동해도 결과가 변하지 않는다.
밀어 올리는 총량과 눌러 내리는 총량이 정확히 같고, 확률을 재분배할 뿐 전체를 통째로 올리거나 내리지 않는다.
이것이 F.cross_entropy가 실제 PyTorch에서 단일 커널인 이유이기도 하다. 수치적으로 안정적이고, backward에 필요한 중간 텐서가 probs 하나뿐이라 메모리도 아끼고, 커널 실행도 8번에서 1번으로 준다.
Deep Learning
· 2026-07-26
-
Andrej Karpathy - Makemore Part 3: Activations & Gradients, BatchNorm
원본 영상 및 자료
Youtube
Github
해당 영상에서는 모델의 학습과정에서 생길 수 있는 Activation/Gradient관련 문제점과, 원활한 학습을 위해 이를 어떻게 해결하는지 소개한다.
1. 이전 내용
이전 영상에서의 MLP코드 위에서 진행된다. block_size는 3으로 동일하고, 글자들의 임베딩 공간 C는 [27, 10]의 shape으로 각 character는 10차원의 공간에서 표현된다.
또한 MLP의 첫 hidden차원이 200으로 확장되었다.
block_size = 3
n_embd = 10 # the dimensionality of the character embedding vectors
n_hidden = 200 # the number of neurons in the hidden layer of the MLP
g = torch.Generator().manual_seed(2147483647) # for reproducibility
C = torch.randn((vocab_size, n_embd), generator=g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator=g)
b1 = torch.randn(n_hidden, generator=g)
W2 = torch.randn((n_hidden, vocab_size), generator=g)
b2 = torch.randn(vocab_size, generator=g)
parameters = [C, W1, W2, b1, b2]
for p in parameters:
p.requires_grad = True
Xb는 각 샘플에서의 context가 담겨있어 [batch_size, 3]의 shape을 가진다. ([0, 0, 5], [0, 5, 13]…)
이를 C[Xb]로 임베딩하면 Xb 안의 숫자들을 행 번호로 간주하여 C에서 그 행들을 그대로 복사해 새 차원을 만들어 추가한다.
최종적으로 emb의 형태는 [batch_size, 3, 10]가 되며, embcat에서 context의 임베딩을 concat하여 [batch_size, 30]이 된다.
for i in range(max_steps):
# minibatch construct
ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)
Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y
# forward pass
emb = C[Xb] # embed the characters into vectors
embcat = emb.view(emb.shape[0], -1) # concatenate the vectors
2. Initialization
logits
현재 코드를 기반으로 학습을 시키며 loss를 출력해보면 다음과 같다.
loss가 학습하며 정상적으로 줄어들고는 있으나, 맨 처음 loss를 보면 27 근처로 굉장히 높은 수치이다.
우리가 사용한 손실함수는 loss = F.cross_entropy(logits, Yb) 이다. batch 크기를 $B$, 정답 character를 $y_i$라 하면
\[\text{loss} = -\frac{1}{B}\sum_{i=1}^{B} \log \frac{\exp(\text{logits}_{i,\,y_i})}{\sum_{j=1}^{27}\exp(\text{logits}_{i,\,j})}\]
즉 logits를 softmax로 확률분포로 바꾼 뒤, 정답 character에 할당된 확률에 $-\log$를 취해 평균낸 값이다.
학습 초기의 상태에는, 어떤 input이 들어오든 모든 character에 대해 동등한 확률을 가지는 것이 자연스럽다. 총 27개의 character가 있으니 초기에는 어떤 character든 1/27 확률로 예측하는 상태여야 한다는 것이다.
이 경우의 loss는
\[-\log\frac{1}{27} = \log 27 \approx 3.2958\]
이 되어야 한다. 하지만 실제 초기 loss는 27 근처이고, 이는 모델이 시작부터 엉뚱한 character에 강한 확신(fake confidence)을 가지고 있다는 뜻이다.
softmax는 입력 전체에 상수를 더해도 결과가 변하지 않으므로, uniform distribution을 만들려면 logits의 각 행이 “전부 0”일 필요는 없고 “전부 같은 값”이기만 하면 된다. 하지만 랜덤 초기화로 값들이 정확히 같아지도록 만들 수는 없으므로, 대신 logits 전체를 0에 가깝게 눌러서 값들 간의 차이를 작게 만드는 방식을 택한다.
logits = h @ W2 + b2 이다. 0에 가까운 값으로 초기화를 해야하므로 randn으로 초기화된 bias(b2)를 0으로 초기화하고, W2 또한 scaling을 통해 작은 값으로 만들어 해결할 수 있다.
W2 = torch.randn((n_hidden, vocab_size), generator=g) * 0.01
b2 = torch.randn(vocab_size, generator=g) * 0
여기서 W2를 정확히 0으로 만들지 않고 0.01을 곱해 아주 작은 값만 남겨두는데, 완전한 대칭 상태를 만들지 않기 위해 약간의 entropy를 남겨두는 것이다.
그 결과 초기 loss가 3.3221로 떨어졌으며, 최종 train/val loss도 함께 떨어졌다. 즉 초기의 부자연스러운 값을 되돌리는 데 학습 cycle이 낭비되지 않고, 처음부터 실제 문제를 optimize하는 데 학습 cycle이 사용된 것이다.
hidden state
logits의 초기값은 안정화가 되었는데, logits를 구성하는 h에 대해서도 살펴볼 필요가 있다. 현재 초기 h의 분포를 확인하면 아래와 같다.
대부분의 값이 -1 또는 1에 분포해있다. 현재 forward pass코드를 보면
# forward pass
emb = C[Xb] # embed the characters into vectors
embcat = emb.view(emb.shape[0], -1) # concatenate the vectors
hpreact = embcat @ W1 + b1 # hidden layer pre-activation
h = torch.tanh(hpreact) # hidden layer
logits = h @ W2 + b2 # output layer
loss = F.cross_entropy(logits, Yb) # loss function
h는 tanh의 출력 결과인데, -1과 1에 몰려있다는 것은 입력인 hpreact가 과하게 양쪽 끝으로 퍼져 있어 tanh가 saturation 되었다는 뜻이다. hpreact의 분포를 확인해보면 대략 -15 ~ 15 범위에 퍼져 있는 것을 확인할 수 있다.
이것이 학습에 문제가 되는 이유는 backward pass에 있다. tanh의 국소 미분은
\[\frac{\partial}{\partial x}\tanh(x) = 1 - \tanh(x)^2\]
이므로, 출력값 $t = \tanh(x)$가 $\pm 1$에 가까울수록 $1 - t^2$가 0에 가까워진다.
위의 코드처럼 tanh의 값이 1이나 -1에 가까울수록 out.grad의 값과 상관없이 gradient를 0에 가깝게 만들어버린다. 즉 tanh가 gradient를 죽여버려서, 그 앞단에 연결된 W1, b1, C 등의 파라미터들이 제대로 된 역전파를 받지 못한다.
더 나아가, 특정 뉴런이 모든 batch 샘플에 대해 saturation 상태라면 그 뉴런은 어떤 입력에도 gradient를 흘려보내지 못하는 dead neuron이 된다. 이 경우 해당 뉴런은 학습 내내 갱신되지 않는다.
이를 해결하기 위해 hpreact 또한 0 주변의 값으로 초기화되도록 해야 한다. 이를 위해 이전과 같이 W1, b1을 scaling하여 완화할 수 있다.
W1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * 0.2
b1 = torch.randn(n_hidden, generator=g) * 0.01
다만 이 0.2라는 값은 분포를 눈으로 확인하며 찾은 임의의 수치이다. 이 scaling 값을 정하는데 비교적 공식화된 방법이 존재한다.
Kaiming He Initialization
아래의 코드처럼, 정규분포 $N(0,1)$을 따르는 x와 w가 있을 때, y = x @ w 라면 y에서 평균은 0 근처로 유지가 되지만, 표준편차가 커진다.
이는 matmul이 fan_in개의 항을 더하는 연산이기 때문이다. 출력의 한 원소는
\[y_j = \sum_{i=1}^{n_{in}} x_i w_{ij}\]
이다. 일반적으로 합의 분산은 각 항의 분산에 공분산 항이 더해지지만, 가중치 $w_{ij}$들이 서로 독립이고 평균이 0이므로 서로 다른 항 사이의 공분산이 0이 되어 분산이 단순히 더해진다.
\[\mathrm{Cov}(x_i w_{ij},\, x_k w_{kj}) = \mathbb{E}[x_i x_k]\cdot\mathbb{E}[w_{ij}w_{kj}] = 0 \quad (i \neq k)\]
각 항의 분산은 $x$와 $w$가 독립이므로 $\mathbb{E}[x^2]\,\mathbb{E}[w^2]$이고, 둘 다 평균이 0이면 이는 $\mathrm{Var}(x)\mathrm{Var}(w)$와 같다. 따라서
\[\mathrm{Var}(y_j) = n_{in}\cdot \mathrm{Var}(x)\cdot \mathrm{Var}(w)\]
가 된다. 즉 표준편차는 $\sqrt{n_{in}}$배만큼 커진다.
w에 scaling을 적용할 때, 1보다 큰 값을 곱하면 표준편차가 더 커지고, 1보다 작은 값을 곱하면 줄어든다. 표준편차를 1 근처로 유지하려면 커진 만큼을 그대로 되돌리면 되므로, 곱해야 할 값은
\[\frac{1}{\sqrt{n_{in}}}\]
이다. 여기서 fan_in이란 w의 input 차원, 즉 하나의 출력 뉴런이 받아들이는 입력의 개수를 말한다.
지금 예제는 bias도 없고 비선형층도 없는 단순한 linear layer여서 간단히 해결했으나, 실제 network는 더 복잡한 구조를 가지고 있을 것이고, 이러한 경우에도 결과값의 분포가 잘 유지되도록 W를 초기화해야 한다.
특히 비선형 함수가 들어오면, 예를 들어 tanh는 입력을 (-1, 1) 범위로 압축하는 함수이므로, 층을 통과할 때마다 분포가 조금씩 줄어든다. 앞서 구한 $1/\sqrt{n_{in}}$만 적용하면 층이 깊어질수록 activation이 0으로 수축해 사라지게 된다. 따라서 이 압축량을 상쇄해줄 보정 계수가 추가로 필요하고, 이를 gain이라 부른다.
이를 정리한 것이 Kaiming He Initialization이다. 원 논문은 ReLU 계열을 다룬 Delving Deep into Rectifiers (He et al., 2015)이며, 가중치의 표준편차를 다음과 같이 정한다.
\[\text{std} = \frac{\text{gain}}{\sqrt{n_{in}}}\]
PyTorch에서는 torch.nn.init.kaiming_normal_로 제공되고, gain 값은 torch.nn.init.calculate_gain(nonlinearity)로 확인할 수 있다.
nonlinearity
gain
Linear / Identity, Conv{1,2,3}D
$1$
Sigmoid
$1$
Tanh
$5/3 \approx 1.667$
ReLU
$\sqrt{2} \approx 1.414$
Leaky ReLU
$\sqrt{2 / (1 + \text{negative_slope}^2)}$
SELU
$3/4$
압축이 강한 함수일수록 gain이 커지는 것을 볼 수 있다. ReLU는 입력의 음수 절반을 통째로 0으로 만들어 분산을 대략 절반으로 줄이므로, 이를 되돌리기 위해 $\sqrt{2}$를 곱한다. tanh 역시 squashing 함수이므로 1보다 큰 $5/3$이 필요하고, 반대로 Linear는 아무것도 압축하지 않으므로 gain이 1이다.
kaiming_normal_에는 mode 인자로 fan_in과 fan_out을 선택할 수 있다. fan_in은 forward pass에서 activation의 분산을 보존하고, fan_out은 backward pass에서 gradient의 분산을 보존한다. 논문에서는 둘 중 어느 쪽을 써도 무방하다고 언급하며, PyTorch의 기본값은 fan_in이다.
우리의 코드에서는 tanh가 사용되었으므로 다음과 같이 scaling을 적용할 수 있다. fan_in은 embcat의 차원인 n_embd * block_size = 30이다.
W1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * (5/3)/((n_embd * block_size)**0.5)
b1 = torch.randn(n_hidden, generator=g) * 0.01
앞서 실험적으로 찾았던 0.2 대신, 이제는 $\frac{5/3}{\sqrt{30}} \approx 0.304$라는 값을 원리에 근거해 얻게 되었다.
기존에는 이러한 initialization 과정이 매우 중요했으나, 오늘날에는 residual connection, normalization layer, 그리고 RMSProp/Adam과 같은 현대적인 optimizer 덕분에 안정성이 많이 확보되어 initialization에 매우 정확한 값을 요구하지는 않는 편이다.
3. BatchNorm
앞서 hpreact의 분포가 적당한 범위에 놓이도록 초기화 시점에 W를 scaling했다. 하지만 이는 어디까지나 간접적인 방법이고, 층이 깊어질수록 이런 계산을 매번 하는 것은 까다롭다.
hpreact가 적절한 분포를 갖기를 원한다면, 초기화를 통해 우회하지 말고 그냥 직접 정규화해버리는 것이 BatchNorm의 시작점이다.
원리
수식은 다음과 같다. 각 미니배치마다 hpreact의 평균과 분산을 구해 이를 바탕으로 정규화한 뒤, scale and shift 과정을 거치는 것이다.
normalize의 epsilon은 분산이 0이 되어 0으로 나누는 경우를 방지하기 위해 추가된 아주 작은 값(PyTorch 기본값 1e-5)이다.
평균과 분산을 batch 차원(dim=0)에 대해 구한다. 200개의 hidden 뉴런 각각에 대해, batch 안의 32개 샘플이 만든 값들의 통계를 계산하여, 결과적으로 평균과 분산은 [1, n_hidden] shape이 된다.
코드로 보면 다음과 같다.
bngain = torch.ones((1, n_hidden))
bnbias = torch.zeros((1, n_hidden))
parameters = [C, W1, W2, b2, bngain, bnbias]
hpreact = bngain * (hpreact - hpreact.mean(0, keepdim=True)) / hpreact.std(0, keepdim=True) + bnbias
정규화만 하면 hpreact는 항상 평균 0, 표준편차 1로 고정되어 버린다. 이는 초기에는 바람직하지만, 학습이 진행되면서 tanh를 saturate시키는 편이 좋다고 판단하더라도 그럴 자유가 없어진다. 그래서 학습 가능한 파라미터 $\gamma$(bngain), $\beta$(bnbias)를 두어 분포를 다시 조정할 수 있게 한다. 학습 초기에는 gain = 1, bias = 0인 상태에서 출발하고, 이후 backprop을 통해 적절한 값으로 조정된다.
이러한 normalization은 보통 linear 혹은 convolution 연산 직후, 비선형성이 추가되는 activation layer 이전에 적용하여 activation의 입력값이 안정되도록 한다. (resnet구조 참고)
문제점
배치 내 샘플들이 서로 얽힌다. 원래 신경망은 하나의 입력에 대해 하나의 출력을 내는 함수여야 하는데, BatchNorm을 넣는 순간 어떤 샘플의 hidden activation이 같은 배치에 우연히 함께 뽑힌 다른 샘플들에 의해 달라진다. 동일한 입력이라도 어떤 batch에 속하느냐에 따라 결과가 바뀌는 상태가 되는 것이다.
그런데 이 버그처럼 보이는 성질이 실제로는 도움이 될 수도 있다. 매 batch마다 평균/분산이 조금씩 다르므로 activation에 일종의 noise가 주입되고, 이것이 data augmentation과 유사한 regularization 효과를 내어 overfitting을 억제한다고 한다.
이 애매한 특징으로 인해 LayerNorm, GroupNorm, InstanceNorm 등이 등장했다.
추론 단계
학습 시에는 batch 통계를 쓰면 되지만, 실제 서비스에서는 샘플 하나만 들어오는 경우가 대부분이다. 샘플이 하나뿐이면 평균과 분산을 구할 수 없다. 따라서 학습이 끝난 뒤 사용할 고정된 평균/표준편차가 필요하다.
가장 단순한 방법은 학습이 끝난 후 전체 training set을 한 번 통과시켜 통계를 계산해두는 것이다.
with torch.no_grad():
emb = C[Xtr]
embcat = emb.view(emb.shape[0], -1)
hpreact = embcat @ W1
bnmean = hpreact.mean(0, keepdim=True)
bnstd = hpreact.std(0, keepdim=True)
하지만 학습 후 별도의 단계를 두는 것은 번거롭다. 그래서 실제 구현에서는 학습 도중에 exponential moving average(EMA)로 통계를 함께 추정해 둔다.
bnmean_running = torch.zeros((1, n_hidden))
bnstd_running = torch.ones((1, n_hidden))
# 학습 루프 안에서
bnmeani = hpreact.mean(0, keepdim=True)
bnstdi = hpreact.std(0, keepdim=True)
hpreact = bngain * (hpreact - bnmeani) / bnstdi + bnbias
with torch.no_grad():
bnmean_running = 0.999 * bnmean_running + 0.001 * bnmeani
bnstd_running = 0.999 * bnstd_running + 0.001 * bnstdi
bnmeani, bnstdi는 이전과 동일하게 배치 내에서의 mean과 std역할을 하고, running값에 매번 일정 비율로 현재 mean과 std값을 누적시킨다.
즉 running 값들은 gradient descent로 학습되는 파라미터가 아니라 따로 갱신하는 값이다. 계산 후 실제 값과 비교해보면 running 추정치가 거의 일치하는 것을 확인할 수 있다.
momentum(위 코드의 0.001, PyTorch 기본값은 0.1)은 batch 크기와 함께 고려해야 한다. batch가 작아 배치별 통계가 요동칠수록 momentum을 작게 두어야 안정된다.
PyTorch BatchNorm의 momentum은 새 관측값 쪽 계수이다: running = (1 - momentum) * running + momentum * new
앞 레이어의 bias가 무의미해진다
hpreact = embcat @ W1 + b1에서 b1은 모든 샘플에 동일한 상수를 더한다. 그런데 BatchNorm은 곧바로 배치 평균을 빼버리므로, 그 상수는 평균에도 똑같이 반영되어 상쇄된다. b1은 결과에 아무 영향을 주지 못하면서 gradient만 낭비하게 된다.
실질적인 bias 역할은 bnbias($\beta$)가 대신하므로, BatchNorm 앞의 linear layer에서는 bias를 아예 제거하는 것이 맞다. (PyTorch에서도 BatchNorm 앞의 nn.Linear나 nn.Conv2d에는 관례적으로 bias=False를 준다.)
hpreact = embcat @ W1 # b1 제거
4. Diagnostics
gain값에 따른 network의 activation/grdient의 변화
linear - tanh의 반복이 5번 쌓인 네트워크에서, 적절한 gain(5/3)을 유지하니 activation의 분포가 안정되고, 1을 사용하였더니 점점 std가 줄어든다. 이는 gradient에서도 마찬가지이다.
data ratio
지금까지는 activation과 gradient의 분포를 봤는데, 학습이 적절한 속도로 진행되고 있는지는 이것만으로 알기 어렵다. gradient의 절대적인 크기는 그 자체로 의미가 없기 때문이다. 파라미터 값이 크다면 큰 gradient도 상대적으로는 미미한 변화이고, 반대로 파라미터가 작다면 작은 gradient도 큰 변화를 일으킬 수 있다.
따라서 봐야 할 것은 한 스텝의 업데이트량이 파라미터 자체의 크기 대비 얼마나 되는가이다.
\[\text{ratio} = \frac{\text{std}(\text{lr} \cdot \nabla_p)}{\text{std}(p)}\]
분자는 이번 스텝에서 실제로 파라미터가 변화하는 양이고, 분모는 파라미터 자체의 스케일이다. 텐서 전체의 대표적인 크기를 재기 위해 std를 사용한다.
학습 루프에서 다음과 같이 기록해둔다.
ud = []
# 학습 루프 안, update 직후
with torch.no_grad():
ud.append([((lr*p.grad).std() / p.data.std()).log10().item() for p in parameters])
log10을 취하는 이유는 값의 범위가 매우 넓기 때문이다. 이렇게 하면 비율 $10^{-3}$이 $-3$으로 표현된다.
plt.figure(figsize=(20, 4))
legends = []
for i, p in enumerate(parameters):
if p.ndim == 2:
plt.plot([ud[j][i] for j in range(len(ud))])
legends.append('param %d' % i)
plt.plot([0, len(ud)], [-3, -3], 'k') # 기준선
plt.legend(legends);
여기서 p.ndim == 2 조건으로 weight matrix만 그리는데, bias나 bngain 같은 1차원 파라미터는 개수가 많아 그래프가 지저분해지기 때문이다.
경험적으로 이 비율은 $10^{-3}$, 즉 로그 스케일에서 $-3$ 근처가 적절하다고 알려져 있다. 한 스텝마다 파라미터가 자기 크기의 약 0.1%씩 변한다는 뜻이다.
$-3$보다 한참 위 ($-2$, $-1$ 등): 한 번의 업데이트가 파라미터를 너무 크게 변화시킴. learning rate를 낮춰야 한다.
$-3$보다 한참 아래 ($-4$, $-5$ 등): 파라미터가 사실상 거의 갱신되지 않고 있다. 이 파라미터는 학습되고 있지 않은 것에 가까우므로 learning rate를 높여야 한다.
학습 초반 스텝은 값이 크게 요동치므로, 안정화된 구간을 기준으로 판단해야 한다.
BatchNorm을 되살리면
앞선 코드에 BatchNorm을 다시 되살린다. 각 Linear 뒤에 BatchNorm1d를 붙인다.
layers = [
Linear(n_embd * block_size, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
Linear( n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
Linear( n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
Linear( n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
Linear( n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
Linear( n_hidden, vocab_size, bias=False), BatchNorm1d(vocab_size),
]
with torch.no_grad():
# last layer: make less confident
layers[-1].gamma *= 0.1
# all other layers: apply gain
for layer in layers[:-1]:
if isinstance(layer, Linear):
layer.weight *= 1.0 # 5/3이 아니라 1.0
gain이 5/3에서 1.0으로 바뀌었다. 즉 Kaiming의 tanh gain을 적용하지 않았다.
마지막 layer를 눌러 초기 confidence를 낮추는 작업이 weight *= 0.1에서 gamma *= 0.1로 바뀌었다. BatchNorm이 앞선 Linear의 weight 스케일을 무시하므로, 출력 스케일을 조절하려면 BatchNorm의 gain을 건드려야 한다.
모든 Linear가 bias=False가 되었다. 앞서 본 대로 BatchNorm 뒤에서는 bias가 상쇄되기 때문이다.
이 상태로 activation 통계를 찍어보면 다음과 같다.
layer 2 (Tanh): mean -0.00, std 0.63, saturated: 2.78%
layer 5 (Tanh): mean +0.00, std 0.64, saturated: 2.56%
layer 8 (Tanh): mean -0.00, std 0.65, saturated: 2.25%
layer 11 (Tanh): mean +0.00, std 0.65, saturated: 1.69%
layer 14 (Tanh): mean +0.00, std 0.65, saturated: 1.88%
gain을 적용하지 않았음에도 층이 깊어져도 std가 0.63~0.65로 유지되고, saturation도 2% 안팎에서 안정적이다. BatchNorm 없이 gain=1.0으로 두었을 때 층마다 activation이 수축해가던 것과 대조적이다.
이는 우연이 아니라 수학적으로 보장되는 성질이다. Linear의 가중치 $W$를 임의의 상수 $a$배 했다고 하자. 그러면 pre-activation도 $a$배가 된다.
\[\text{hpreact} \rightarrow a \cdot \text{hpreact}\]
그런데 BatchNorm은 이 값의 배치 평균과 표준편차를 구해 나눈다. 평균도 $a\mu$, 표준편차도 $a\sigma$가 되므로
\[\frac{a x - a\mu}{a\sigma} = \frac{x - \mu}{\sigma}\]
로 $a$가 정확히 소거된다. 즉 BatchNorm 앞에 있는 Linear layer의 weight 스케일은 forward pass 결과에 아무런 영향을 주지 못한다. 앞서 계산했던 gain과 $1/\sqrt{n_{in}}$은, BatchNorm이 붙는 순간 출력 분포에 관해서는 무의미해지는 것이다.
이것이 BatchNorm이 널리 쓰이게 된 가장 큰 실용적 이유다. 층이 수십, 수백 개로 깊어지면 각 층의 분포를 초기화만으로 맞추는 것은 사실상 불가능한데, BatchNorm은 그 부담을 통째로 없애준다. 초기화가 다소 어긋나 있어도 네트워크는 여전히 학습된다.
그래도 완전히 무관한 것은 아니다 (보충 설명)
영상 내용 외에 추가로 정리.
BatchNorm을 하나의 함수 $f$로 보먄, 앞서 확인한 스케일 불변성은 다음과 같이 쓸 수 있다.
\[f(az) = f(z) \qquad (a > 0)\]
양변을 $z$로 미분한다. 좌변은 연쇄법칙에 의해 안쪽 미분 $a$가 곱해지므로
\[a \cdot f'(az) = f'(z)
\quad\Longrightarrow\quad
f'(az) = \frac{1}{a} f'(z)\]
forward는 그대로인데 미분값은 $1/a$배로 줄어든다.
서로 다른 두 개의 신경망이 있다고 하자.
네트워크 A: 가중치가 $W$
네트워크 B: 가중치가 $W’ = aW$ ($W$의 모든 원소에 똑같은 숫자 $a$를 곱한 것 (예: $a=2$면 전부 2배))
두 네트워크는 가중치 숫자가 다르다. 그런데도 BatchNorm을 통과하면 출력이 같다는 것이 앞서 확인한 내용이었다.
그러면 “출력이 같으니 학습도 똑같이 되는가?” 답은 “아니다”.
역전파의 출발점은 같다
역전파는 출력 쪽에서 시작해 입력 쪽으로 거슬러 올라간다.
... → z = Wx → BatchNorm → y → 손실 L
↑
여기서 역전파 시작
$L$은 $y$와 정답만으로 계산된다. 그런데 A와 B는 $y$가 완전히 똑같으므로
손실 $L$도 똑같다
그 손실을 $y$로 미분한 값 $\partial L / \partial y$도 똑같다
$\partial L/\partial y$는 “출력을 어느 방향으로 얼마나 밀어야 손실이 줄어드는가”를 나타내는 신호다. 두 네트워크가 완전히 같은 출력을 냈으니, 받는 신호도 똑같다.
차이는 이 신호가 BatchNorm을 거슬러 통과하는 순간부터 생긴다.
BatchNorm을 통과하면 $1/a$배가 된다
이제 그 신호를 $y$에서 $z$로 옮겨야 한다. 연쇄법칙에 따라
\[\underbrace{\frac{\partial L}{\partial z}}_{\text{구하려는 것}} = \underbrace{\frac{\partial L}{\partial y}}_{\text{A, B 동일}} \times \underbrace{\frac{\partial y}{\partial z}}_{\text{여기서 차이}}\]
앞항은 방금 봤듯이 같다. 그러니 $\partial y / \partial z$, 즉 BatchNorm 자체의 미분값만 비교하면 된다.
$\partial y/\partial z$가 뜻하는 것은 “BatchNorm에 들어간 값을 변화시키면 출력이 얼마나 변하는가”이다.
A
B ($a=2$)
$z$의 표준편차
1
2
BN이 나누는 값
1
2
$z$가 $0.1$ 변화
출력이 $0.1/1 = 0.1$ 변함
출력이 $0.1/2 = 0.05$ 변함
B는 z를 같은 크기를 변화시켜도 출력이 절반밖에 안 움직인다. BatchNorm이 나누는 표준편차가 2배로 커졌기 때문이다.
“입력 변화에 덜 반응한다”가 곧 “미분값이 작다”이다. 그래서 $\partial y/\partial z$가 $1/a$배가 되고, 결과적으로
\[\frac{\partial L}{\partial z'} = \frac{1}{a}\cdot\frac{\partial L}{\partial z}\]
가 된다.
$W$를 크게 잡을수록 뒤에서 오는 학습 신호가 앞으로 전달될 때 더 많이 감쇄된다.
$z = Wx$이므로 $W$에 대한 gradient는
\[\frac{\partial L}{\partial W} = x^\top \frac{\partial L}{\partial z}\]
인데, 여기서 $x$는 A와 B가 완전히 같은 입력을 쓴다. 그러니 앞에서 $1/a$배가 된 것이 그대로 유지된다.
\[\frac{\partial L}{\partial W'} = \frac{1}{a}\cdot\frac{\partial L}{\partial W}\]
가중치는 $a$배 크고, gradient는 $1/a$배 작다. 이 둘이 합쳐져 결국 학습이 $1/a^2$만큼 느려진다.
실효 learning rate는 $1/a^2$배가 된다
앞 절의 data ratio를 다시 가져오자.
\[\text{ratio} = \frac{\|\,\text{lr}\cdot \nabla_W\,\|}{\|W\|}\]
$W$를 $a$배 했을 때 분자와 분모가 각각 어떻게 변하는지 보면:
스케일
변화
분모 (파라미터 크기)
$|aW| = a|W|$
$a$배 증가
분자 (업데이트 크기)
$\text{lr}\cdot\frac{1}{a}|\nabla_W|$
$1/a$배 감소
\[\text{ratio}' = \frac{1}{a^2}\cdot\text{ratio}\]
두 효과가 같은 방향으로 겹치기 때문에 $1/a$가 아니라 $1/a^2$이다. 초기화 스케일을 2배로 잡으면 그 layer는 실효적으로 learning rate가 1/4로 줄어든 채 학습되는 셈이고, 4배로 잡으면 1/16이 된다. forward 출력은 완전히 동일한데도 말이다.
직접 확인해보기
import torch
torch.manual_seed(0)
x = torch.randn(32, 30)
W = torch.randn(30, 100) / 30**0.5
target = torch.randn(32, 100)
bn = torch.nn.BatchNorm1d(100, affine=False)
def run(scale):
Wa = (scale * W).clone().requires_grad_(True)
y = bn(x @ Wa)
loss = 0.5 * ((y - target)**2).sum()
loss.backward()
return y.detach(), Wa.grad.norm().item(), Wa.detach().norm().item()
y0, g0, w0 = run(1.0)
for a in [0.25, 4.0, 16.0]:
y, g, w = run(a)
print(f"a={a:5.2f} forward 최대오차={(y-y0).abs().max():.2e}"
f" |grad|비={g/g0:8.4f} (1/a={1/a:.4f})"
f" ratio비={(g/w)/(g0/w0):9.5f} (1/a²={1/a**2:.5f})")
a= 0.25 forward 최대오차=5.6e-04 |grad|비= 3.9996 (1/a=4.0000) ratio비= 15.99830 (1/a²=16.00000)
a= 4.00 forward 최대오차=3.5e-05 |grad|비= 0.2500 (1/a=0.2500) ratio비= 0.06250 (1/a²=0.06250)
a=16.00 forward 최대오차=3.7e-05 |grad|비= 0.0625 (1/a=0.0625) ratio비= 0.00391 (1/a²=0.00391)
forward 출력은 (오차 $10^{-4}$ 이하로) 완전히 동일한데, gradient 크기는 정확히 $1/a$, update:data ratio는 정확히 $1/a^2$로 변한다.
forward 오차가 정확히 0이 아닌 것은 분모의 $\epsilon$ 때문이다. $\sqrt{a^2\sigma^2 + \epsilon}$는 $a\sqrt{\sigma^2+\epsilon}$와 완전히 같지 않아서, $\epsilon$이 스케일 불변성을 아주 미세하게 깨뜨린다.
정리
BatchNorm은 forward 출력에 대해서만 초기화 스케일을 완전히 지워준다.
스케일에 따라 학습 속도에 대해서는 실효 learning rate를 $1/a^2$ 비율로 변함.
BatchNorm은 W의 안정화된 분포를 만들어주지만, 각 layer가 올바른 속도로 학습되는지는 data ratio로 확인해야 한다.
Deep Learning
· 2026-07-23
-
Andrej Karpathy - Makemore Part 2: MLP
원본 영상 및 자료
Youtube
Github
해당 영상에서는 “A Neural Probabilistic Language Model”(링크) 논문에서의 MLP를 이용하여 Character단위의 Trigram 모델을 구현한다.
위의 사진과 비슷한 구조의 MLP를 구현할 것이며, 논문에서는 가까운 단어 3개를 인풋으로 사용하지만 영상에서는 글자를 단위로 구현하며, 각 글자에 대한 임베딩을 concat하여 input으로 사용한다.
1. 데이터 로드 및 Vocabulary 구축
words = open('names.txt', 'r').read().splitlines()
chars = sorted(list(set(''.join(words))))
stoi = {s:i+1 for i,s in enumerate(chars)}
stoi['.'] = 0
itos = {i:s for s,i in stoi.items()}
''.join(words): words에 있는 모든 단어를 공백 없이 하나의 긴 문자열로 합친다.
set(…): 문자열을 set 자료형으로 만들어 중복된 character를 모두 제거
list(…) & sorted(…): 순서가 없는 set을 다시 리스트로 변환, a부터 z까지 오름차순으로 정렬한다.
이후 '.'을 0번 인덱스에 할당하고 stoi, itos를 생성.
2. 데이터셋 구축 (Sliding Window)
block_size = 3 # context length
X, Y = [], []
for w in words:
context = [0] * block_size
for ch in w + '.':
ix = stoi[ch]
X.append(context)
Y.append(ix)
context = context[1:] + [ix] # crop and append
X = torch.tensor(X)
Y = torch.tensor(Y)
각 단어 w에 대해, 처음 context는 모두 block_size만큼 . 이 반복된 형태를 가진다.
X에는 context를 할당하고, Y에는 현재 위치의 글자 인덱스를 할당한다. 이후 context = context[1:] + [ix]를 통해 맨 앞 글자는 잘라내고, 현재 [ix]를 추가하여 window를 한 칸씩 옆으로 이동시킨다.
3. Embedding Lookup Table
각 글자를 모델에 넣기 위해 인덱스에 해당하는 글자를 벡터로 변환하는 임베딩이 필요하다. 영상에서는 처음에 글자들을 모두 2차원의 벡터 공간에 담는다.
C = torch.randn((27, 2))
emb = C[X]
emb.shape # torch.Size([228146, 3, 2])
C는 글자들의 임베딩 공간이다. 알파벳 26개 + 점(.) 1개로 27개의 글자 각각의 2차원 벡터 표현을 랜덤으로 선언한 것이다.
X는 2번 과정에서 만든 텐서이다. 각 샘플에서의 context가 담겨있으므로 [데이터개수, 3]의 shape을 가진다. ([0, 0, 5], [0, 5, 13]…)
여기서 emb = C[X]를 하게 되면, X 안에 적힌 숫자들을 행 번호로 간주하여 C에서 그 행들을 그대로 복사해 새 차원을 만들어 추가한다.
예를 들어 X 안에 [0, 5, 13]이 들어있다면 C의 0번째, 5번째, 13번째 행을 가져와 마지막 차원에 추가한다. 인덱스 하나가 들어있던 자리가 2차원 벡터로 바뀐 것이다.
최종적으로 emb의 형태는 [데이터개수, 3, 2]가 되며, 이러한 방식은 적은 연산량으로 효과적으로 범주형 변수를 벡터로 변환하는 유용한 인덱싱 방식이다.
4. 모델 구조 및 Loss (Cross Entropy)
W1 = torch.randn((6, 100))
b1 = torch.randn(100)
h = torch.tanh(emb.view(-1, 6) @ W1 + b1)
W2 = torch.randn((100, 27))
b2 = torch.randn(27)
logits = h @ W2 + b2
emb.view(-1, 6)을 통해 3차원 형태였던 임베딩을 2차원(글자 3개 $\times$ 2차원 = 6)으로 바꾼 뒤, Linear Layer 두 개와 tanh 활성화 함수를 통과시킨다. logits의 shape은 [데이터개수, 27]로, 각각의 sample에서 다음 character의 확률을 예측해내야 한다. 이를 위해 softmax를 사용한다.
counts = logits.exp()
prob = counts / counts.sum(1, keepdims=True)
이후 모델의 예측 확률(prob)과 실제 정답(Y) 간의 오차를 구하기 위해 Cross Entropy를 사용한다.
loss = -prob[torch.arange({데이터 개수}), Y].log().mean()
torch.arange({데이터 개수})를 통해 행 번호를 지정하고, Y를 통해 열 번호를 지정하여 정답 확률을 가져온다. 그 후 로그를 씌워 평균을 낸 후 음수를 취한다. (Negative Log Likelihood, NLL)
위 방식은 수학적으로는 완벽하지만, 파이썬 코드상으로는 치명적인 약점이 있다. 만약 logits에 100, 1000처럼 큰 값이 들어갈 경우, logits.exp()를 계산하는 과정에서 값이 무한대(inf)로 폭발하여 정상적인 계산이 불가능해진다(NaN 발생).
따라서 실제 Softmax 연산을 구현할 때는 수치적 안정성을 위해 배열 내 최댓값(max)을 모든 값에서 빼주는 방식으로 해결한다. 지수 함수의 특성 상 분모와 분자에 동일한 값을 곱해 상쇄시키는 원리다.
\[p_i = \frac{e^{x_i}}{\sum_{j} e^{x_j}} = \frac{e^{x_i} \cdot e^{-M}}{\sum_{j} (e^{x_j} \cdot e^{-M})} = \frac{e^{x_i - M}}{\sum_{j} e^{x_j - M}} \quad (\text{단, } M = \max(x))\]
또한, 파이토치에서 제공하는 내장 함수인 F.cross_entropy를 사용하면 이 모든 과정을 한 번에 처리할 수 있다. 안정성 처리가 된 Softmax 연산이 포함되어 있는 손실함수이다(Softmax + Log + NLLLoss).
logits = h @ W2 + b2
loss = F.cross_entropy(logits, Ytr[ix])
5. 최적의 학습률 찾기
lre = torch.linspace(-3, 0, 1000)
lrs = 10**lre
# 학습 루프 내부
...
lr = lrs[i]
for p in parameters:
p.data += -lr * p.grad
학습 루프를 돌면서 학습률을 작은 값($10^{-3}$)부터 큰 값($10^0$)까지 서서히 증가시키며 Loss의 변화를 추적하는 기법이다.
미니배치가 매번 바뀌기 때문에 Loss가 위아래로 흔들리는 노이즈는 분명 발생한다. 하지만 학습률이 변하면서 Loss에 미치는 영향력이 훨씬 커서, 그래프를 그려보면 Loss가 가장 가파르게 떨어지는 구간이 눈에 띄게 나타난다. 그 가파른 구간의 학습률(영상에서는 0.1 부근)을 최적의 초기 학습률로 채택한다.
Deep Learning
· 2026-07-18
hemisus
